from keras.preprocessing.image import load_img, img_to_array
import os
import numpy as np


class PreProcess(object):
    def __init__(self):
        pass

    @staticmethod
    def create_training_data(train_src_path, mask_src_path, dst_path):
        train_data, labels = [], []
        count = 0
        for img in os.listdir(mask_src_path):
            if img.endswith("png") or img.endswith("jpg"):
                train_img = load_img("{0}/{1}".format(train_src_path, img), grayscale=True)
                msk_img = load_img("{0}/{1}".format(mask_src_path, img), grayscale=True)

                train_data.append(img_to_array(train_img))
                labels.append(img_to_array(msk_img))

                count += 1
                if count % 100 == 99:
                    print("count:{}".format(count))

        train_data, labels = np.array(train_data), np.array(labels)
        train_data.astype("float32")
        labels.astype("float32")
        train_data /= 255
        labels[labels > 0] = 1

        np.save("{}/train.npy".format(dst_path), train_data)
        np.save("{}/mask.npy".format(dst_path), labels)

    @staticmethod
    def get_training_data(dst_path):
        train_data = np.load("{}/train.npy".format(dst_path))
        mask_data = np.load("{}/mask.npy".format(dst_path))

        return train_data, mask_data

    @staticmethod
    def create_test_data(test_src_path, dst_path):
        test_data = []

        count = 0
        for img in os.listdir(test_src_path):
            if img.endswith("png") or img.endswith("jpg"):
                test_img = load_img("{0}/{1}".format(test_src_path, img), grayscale=True)

                test_data.append(img_to_array(test_img))

                count += 1
                if count % 100 == 99:
                    print("count:{}".format(count))

        test_data = np.array(test_data)
        test_data.astype("float32")
        test_data /= 255

        np.save("{}/test.npy".format(dst_path), test_data)

    @staticmethod
    def get_test_data(dst_path):
        test_data = np.load("{}/test.npy".format(dst_path))

        return test_data

    @staticmethod
    def reshape(dir, shape=(512, 512)):
        from PIL import Image
        for img_name in os.listdir(dir):
            img = Image.open("{0}/{1}".format(dir, img_name))
            img = img.resize(shape, Image.ANTIALIAS)
            img.save("{0}/{1}".format(dir, img_name))




if __name__ == "__main__":
    # PreProcess.reshape("../test")
    # PreProcess.create_test_data("../test", "../train/test")
    PreProcess.create_training_data("../train/image", "../train/mask", "../train/dest")
    # PreProcess.get_training_data("../train/dest")
    # print(PreProcess.get_test_data("../train/test").shape)